import torch
from torch import nn
from d2l import torch as d2l
import matplotlib.pyplot as plt
# 读取训练和测试数据集并设置批量大小
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

# 初始化模型参数，单层隐藏层的MLP，每一个层要初始化对应的权重和偏置
num_inputs, num_outputs, num_hiddens = 784, 10, 256
W1 = nn.Parameter(torch.randn(num_inputs, num_hiddens)*0.01, requires_grad=True)
b1 = nn.Parameter(torch.zeros(num_hiddens), requires_grad=True)
W2 = nn.Parameter(torch.randn(num_hiddens, num_outputs)*0.01, requires_grad=True)
b2 = nn.Parameter(torch.zeros(num_outputs), requires_grad=True)
params = [W1, b2, W2, b2]

# 激活函数实现，手动实现不直接使用内置的relu函数
def relu(X):
    a = torch.zeros_like(X) # 直接后面加_like实现变量形状一致
    return torch.max(a, X)  # 注意这里要用torch的max函数！！！

# 模型实现，记得将输入图像展位向量
def net(X):
    X = X.reshape([-1, num_inputs])
    H = relu(X@W1 + b1) # 这里的@代表的是矩阵乘法
    return (H@W2 + b2)

# 损失函数选择交叉熵
loss = nn.CrossEntropyLoss(reduction='none')

# 训练,调用softmax中的训练包，num_epochs设置为10，lr学习率设置为0.1，设置优化算法梯度下降
num_epochs, lr = 10, 0.1
updater = torch.optim.SGD(params, lr=lr)
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, updater)
# d2l.plt.show() 如果要打印训练损失图，控制台加这个方法即可，全部一起执行很奇怪没有曲线
# 使用学习到的模型进行预测
d2l.predict_ch3(net, test_iter)




